package com.handwriting.backend.service;

import com.baomidou.mybatisplus.core.conditions.query.QueryWrapper;
import com.handwriting.backend.Config;
import com.handwriting.backend.api.PythonAPI;
import com.handwriting.backend.entity.Model;
import com.handwriting.backend.exceptions.ModelDefinedException;
import com.handwriting.backend.mapper.ModelMapper;
import com.handwriting.backend.message.PyResponse;
import com.handwriting.backend.message.TrainParams;
import com.handwriting.backend.message.TrainingProgress;
import com.handwriting.backend.utils.JsonUtils;
import org.apache.ibatis.javassist.NotFoundException;
import org.springframework.http.client.MultipartBodyBuilder;
import org.springframework.stereotype.Component;
import org.springframework.web.multipart.MultipartFile;

import java.io.IOException;
import java.net.URI;
import java.util.List;

@Component
public class ModelService {
    private final ModelMapper mapper;
    private final Config config;
    private final Object o = new Object();
    private final PythonAPI pythonAPI;
    private final JsonUtils json;

    public ModelService(ModelMapper mapper, Config config, PythonAPI pythonAPI, JsonUtils json) {
        this.mapper = mapper;
        this.config = config;
        this.pythonAPI = pythonAPI;
        this.json = json;
    }

    public List<Model> getAllModelInfo() {
        List<Model> list = mapper.selectList(null);
        for (Model m : list) {
            m.setSelected(config.value().getModelUsedId() == m.getId());
            m.setTestRatio(1 - m.getTrainRatio() - m.getTestRatio());
        }
        return list;
    }

    public void selectModel(int id) throws NotFoundException, IOException {
        long cnt = mapper.selectCount(new QueryWrapper<Model>().eq("id", id));
        if (cnt <= 0) {
            throw new NotFoundException("Could not find the model id " + id + ".");
        }
        synchronized (o) {
            config.value().setModelUsedId(id);
            config.save();
        }
    }

    public void deleteModel(int id) throws IOException {
        List<Model> lst = mapper.selectList(null);
        if (config.value().getModelUsedId() == id) {
            for (Model m : lst) {
                if (m.getId() != id) {
                    synchronized (o) {
                        config.value().setModelUsedId(id);
                        config.save();
                    }
                }
            }
        }
        mapper.deleteById(id);
    }

    public void trainModel(TrainParams params) {
//        model_name=train_params["modelName"],
//                dataset=train_params["dataset"],
//                num_epochs=train_params["epochs"],
//                batch_size=train_params["batchSize"],
//                learning_rate=train_params["learningRate"],
//                optimizer=train_params["optimizer"],
//                train_ratio=train_params["trainRatio"],
//                val_ratio=train_params["valRatio"],
//                test_ratio=train_params["testRatio"]
//        String modelName = Utils.genModelName();
//        params.setModelName(modelName);
        List<Model> models = mapper.selectList(new QueryWrapper<Model>().eq("name", params.getModelName()));
        if (!models.isEmpty())
            throw new ModelDefinedException("Model has already defined.");
        String responseJson = pythonAPI.post(json.toJson(params), URI.create("http://localhost:5000/model/train/start"));
        PyResponse response = json.fromJson(responseJson, PyResponse.class);
        if (response.error()) {
            throw new RuntimeException("Error starting train.");
        }
//        return modelName;
    }

    public TrainingProgress getTrainingProgress(String modelName) {
        String responseJson = pythonAPI.get(URI.create("http://localhost:5000/model/train/progress/" + modelName));
        System.out.println(responseJson);
        PyResponse response = json.fromJson(responseJson, PyResponse.class);
        return json.fromJson(response.getMsg(), TrainingProgress.class);
    }

    public int testModel(MultipartFile image) {
        int modelId = config.value().getModelUsedId();
        Model model = mapper.selectList(new QueryWrapper<Model>().eq("id", modelId)).get(0);
        String modelName = model.getName();
        MultipartBodyBuilder builder = new MultipartBodyBuilder();
        builder.part("modelName", modelName);
        builder.part("image", image.getResource());
        String result = pythonAPI.postForm(builder, URI.create("http://localhost:5000/model/test"));
        System.out.println(result);
        model.setUsedCount(model.getUsedCount() + 1);
        mapper.updateById(model);
        if (result == null || result.isBlank()) {
            return -1;
        }
        return Integer.parseInt(result);
    }

    public void feedback() {
        int modelId = config.value().getModelUsedId();
        Model model = mapper.selectList(new QueryWrapper<Model>().eq("id", modelId)).get(0);
        model.setErrorCount(model.getErrorCount() + 1);
        mapper.updateById(model);
    }
}
